package com.by.rpc.config;

import com.by.rpc.dto.RpcRequest;
import com.by.rpc.dto.RpcResponse;
import com.by.rpc.dto.ServiceInstance;
import com.by.rpc.util.ServiceParameter;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.Serializer;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.reflections.Reflections;
import org.reflections.scanners.SubTypesScanner;
import org.reflections.util.ClasspathHelper;
import org.reflections.util.ConfigurationBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.lang.reflect.Modifier;
import java.util.Set;
import java.util.UUID;

/**
 * 自定义 UUID 序列化器，用于解决 Java 9+ 的模块化反射问题
 */
class UUIDSerializer extends Serializer<UUID> {
    @Override
    public void write(Kryo kryo, Output output, UUID object) {
        output.writeLong(object.getMostSignificantBits());
        output.writeLong(object.getLeastSignificantBits());
    }

    @Override
    public UUID read(Kryo kryo, Input input, Class<? extends UUID> type) {
        long mostSigBits = input.readLong();
        long leastSigBits = input.readLong();
        return new UUID(mostSigBits, leastSigBits);
    }
}

/**
 * 支持一次性自动扫描并注册VO类的Kryo序列化工具类
 */
public final class KryoSerializer {

    private static final Logger logger = LoggerFactory.getLogger(KryoSerializer.class);

    // 定义需要扫描的包路径
    private static final String[] SCAN_PACKAGES = {
            "com.by.order.api.vo",
            "com.by.user.api.vo",
            "com.by.product.api.vo",
            "com.by.rpc.dto"
    };

    // 静态变量：在类加载时进行一次性、全局的类扫描并缓存结果
    private static final Set<Class<?>> SCANNED_CLASSES;

    static {
        long startTime = System.currentTimeMillis();
        // 执行一次性扫描
        SCANNED_CLASSES = scanClasses(SCAN_PACKAGES);
        long endTime = System.currentTimeMillis();
        logger.info("Kryo registration scan completed in {} ms. Found {} classes to register.",
                (endTime - startTime), SCANNED_CLASSES.size());
    }

    // 线程隔离的Kryo实例，每个线程都有一个独立的 Kryo 对象
    private static final ThreadLocal<Kryo> kryoThreadLocal = ThreadLocal.withInitial(() -> {
        Kryo kryo = new Kryo();
        // 开启严格模式，未注册的类会抛出异常
        kryo.setRegistrationRequired(true);

        // 1. 注册基础框架类和常用 Java 类
        registerBaseClasses(kryo);

        // 2. 从缓存中获取扫描结果并进行注册，不再重复执行反射扫描
        for (Class<?> clazz : SCANNED_CLASSES) {
            // 过滤接口和抽象类
            if (!clazz.isInterface() && !Modifier.isAbstract(clazz.getModifiers())) {
                kryo.register(clazz);
            }
        }

        return kryo;
    });

    /**
     * 注册基础框架类和常用Java类
     */
    private static void registerBaseClasses(Kryo kryo) {
        kryo.register(RpcRequest.class);
        kryo.register(RpcResponse.class);
        kryo.register(ServiceInstance.class);
        kryo.register(ServiceParameter.class);
        kryo.register(RpcResponse.ResponseStatus.class);

        // 注册常用集合和数组类型
        kryo.register(Object[].class);
        kryo.register(String[].class);
        kryo.register(Class[].class);
        kryo.register(byte[].class);
        kryo.register(int[].class);
        kryo.register(boolean[].class);
        kryo.register(java.util.ArrayList.class);
        kryo.register(java.util.HashMap.class);
        kryo.register(java.util.HashSet.class);
        kryo.register(java.util.Date.class);
        kryo.register(java.math.BigDecimal.class);

        // Kryo 对集合的特殊处理，注册其实现类
        kryo.register(java.util.Collections.singletonList(null).getClass());
        kryo.register(java.util.Arrays.asList().getClass());

        // 使用自定义序列化器注册 UUID，解决 JPMS 反射问题
        kryo.register(UUID.class, new UUIDSerializer());
    }

    /**
     * 扫描指定包下的所有类
     * @param basePackages 基础包路径
     * @return 类集合
     */
    private static Set<Class<?>> scanClasses(String... basePackages) {
        ConfigurationBuilder configBuilder = new ConfigurationBuilder()
                .forPackages(basePackages)
                .setScanners(new SubTypesScanner(false));

        for (String basePackage : basePackages) {
            configBuilder.addUrls(ClasspathHelper.forPackage(basePackage));
        }

        Reflections reflections = new Reflections(configBuilder);
        return reflections.getSubTypesOf(Object.class);
    }

    /**
     * 序列化对象
     */
    public static <T> byte[] serialize(T obj) {
        Kryo kryo = kryoThreadLocal.get();
        try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
             Output output = new Output(byteArrayOutputStream)) {
            kryo.writeClassAndObject(output, obj);
            output.flush();
            return byteArrayOutputStream.toByteArray();
        } catch (Exception e) {
            logger.error("Kryo serialize failed", e);
            // 捕获异常后，清理当前线程的Kryo实例，避免实例污染
            kryoThreadLocal.remove();
            throw new RuntimeException("Kryo serialize failed", e);
        }
    }

    /**
     * 反序列化对象
     */
    @SuppressWarnings("unchecked")
    public static <T> T deserialize(byte[] bytes) {
        Kryo kryo = kryoThreadLocal.get();
        try (ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes);
             Input input = new Input(byteArrayInputStream)) {
            Object obj = kryo.readClassAndObject(input);
            return (T) obj;
        } catch (Exception e) {
            logger.error("Kryo deserialize failed", e);
            kryoThreadLocal.remove();
            throw new RuntimeException("Kryo deserialize failed", e);
        }
    }
}